package com.xz.xzutils.module.xzmybatislogplus.socket;

import java.io.DataInputStream;
import java.io.DataOutputStream;
import java.io.IOException;
import java.net.ServerSocket;
import java.net.Socket;
import java.nio.charset.StandardCharsets;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * <p>
 *
 * </p>
 *
 * @author DONG.WANG17@HAND-CHINA.COM
 * @since 2025/3/22-03-22 16:11
 */
public class SocketServer implements SocketEndpoint {
    private final SocketConfig config;
    private ServerSocket serverSocket;
    private final ExecutorService executor = Executors.newCachedThreadPool();
    private final CopyOnWriteArrayList<ClientConnection> clients = new CopyOnWriteArrayList<>();
    private final AtomicBoolean isRunning = new AtomicBoolean(false);

    public SocketServer(SocketConfig config) {
        this.config = config;
        // 关闭钩子
        Runtime.getRuntime().addShutdownHook(new Thread(this::stop));
    }

    @Override
    public void start() {
        if (isRunning.get()) return;

        try {
            serverSocket = new ServerSocket(config.getPort());
            serverSocket.setReuseAddress(true);
            isRunning.set(true);
            executor.execute(this::acceptConnections);
            notifyStatus(ConnectionStatus.CONNECTED);
            config.getCallback().onStart();
        } catch (IOException e) {
            handleError(e);
        }
    }

    private void acceptConnections() {
        while (isRunning.get()) {
            try {
                Socket clientSocket = serverSocket.accept();
                System.out.println("[Server] 客户端连接: " + clientSocket.getInetAddress());
                ClientConnection connection = new ClientConnection(clientSocket);
                clients.add(connection);
                executor.execute(connection);
            } catch (IOException e) {
                if (!serverSocket.isClosed()) handleError(e);
            }
        }
    }

    @Override
    public void sendData(String message) {
        clients.forEach(conn -> {
            if (conn.connected.get()) {
                conn.send(message);
            } else {
                clients.remove(conn);
            }
        });
    }

    @Override
    public void stop() {
        isRunning.set(false);
        try {
            serverSocket.close();
            executor.shutdownNow();
            clients.forEach(ClientConnection::close);
            notifyStatus(ConnectionStatus.DISCONNECTED);
        } catch (IOException e) {
            handleError(e);
        }
    }

    @Override
    public boolean isRunning() {
        return isRunning.get();
    }

    private class ClientConnection implements Runnable {
        private final Socket socket;
        private final DataInputStream input;
        private final DataOutputStream output;
        private final AtomicBoolean connected = new AtomicBoolean(true);
        private final BlockingQueue<String> messageQueue = new LinkedBlockingQueue<>();

        ClientConnection(Socket socket) throws IOException {
            this.socket = socket;
            this.input = new DataInputStream(socket.getInputStream());
            this.output = new DataOutputStream(socket.getOutputStream());
        }

        @Override
        public void run() {
            try {
                // 启动发送线程
                Executors.newSingleThreadExecutor().execute(() -> {
                    while (connected.get()) {
                        try {
                            String message = messageQueue.take();
                            byte[] data = message.getBytes(StandardCharsets.UTF_8);
                            output.writeInt(data.length);
                            output.write(data);
                            output.flush();
                        } catch (InterruptedException e) {
                            Thread.currentThread().interrupt();
                            break;
                        } catch (IOException e) {
                            close();
                            handleError(e);
                            break;
                        }
                    }
                });
                // 接收消息
                while (connected.get()) {
                    int length = input.readInt();
                    byte[] buffer = new byte[length];
                    input.readFully(buffer);
                    String message = new String(buffer, StandardCharsets.UTF_8);

                    if (config.getHeartbeatTag().equals(message) && !config.getReceiveHeartbeatTag()) {
                        continue;
                    }
                    config.getCallback().onMessage(message);
                    final String response = config.getCallback().onMessageAndResponse(message);
                    if (response != null) {
                        // 回复消息
                        send(response);
                    }
                }
            } catch (IOException e) {
                close();
                handleError(e);
            }
        }

        void send(String message) {
            messageQueue.offer(message);
        }

        void close() {
            connected.set(false);
            try {
                socket.close();
                clients.remove(this);
                config.getCallback().onClose();
            } catch (IOException e) {
                handleError(e);
            }
        }
    }

    private void handleError(Exception e) {
        if (config.getCallback() != null) {
            config.getCallback().onError(e);
        }
    }

    private void notifyStatus(ConnectionStatus status) {
        if (config.getCallback() != null) {
            config.getCallback().onStatusChange(status);
        }
    }
}